/*
 * Javlov - a Java toolkit for reinforcement learning with multi-agent support.
 * 
 * Copyright (c) 2009 Matthijs Snel
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

//import gnu.trove.TObjectDoubleHashMap;

/**
 * State-value function for use in discrete state spaces.
 * 
 * @author Matthijs Snel
 */
@SuppressWarnings("unchecked")
public abstract class TabularValueFunction implements ValueFunction {

	private static final long serialVersionUID = -3403768427883527887L;

	/**
	 * Use this factory method if the size of the state space is unknown. Will result in slower
	 * performance than when using {@link #getInstance(int)}.
	 * 
	 * If the size of the state space is unknown because the space is continuous, you probably
	 * shouldn't be using this class, but one of the classes that implement function approximators /
	 * generalisations over states for the value function (yet to be implemented).
	 * 
	 * The use case for this method is one in which the state space is discrete and of 
	 * unknown size (e.g. because the environment is dynamic and hence the size of the state space
	 * is hard to predict), yet finite and small enough for it to make sense to be stored in a
	 * table (as a rough indication I'd say less than a million states).
	 * 
	 * This method returns a value function for TD(0).
	 */
	public static TabularValueFunction getInstance() {
		return getInstance(0);
	}
	
	/**
	 * Use this factory method if the size of the state space is unknown. Will result in slower
	 * performance than when using {@link #getInstance(int)}.
	 * 
	 * If the size of the state space is unknown because the space is continuous, you probably
	 * shouldn't be using this class, but one of the classes that implement function approximators /
	 * generalisations over states for the value function (yet to be implemented).
	 * 
	 * The use case for this method is one in which the state space is discrete and of 
	 * unknown size (e.g. because the environment is dynamic and hence the size of the state space
	 * is hard to predict), yet finite and small enough for it to make sense to be stored in a
	 * table (as a rough indication I'd say less than a million states).
	 * 
	 * @param traceDecay the decay parameter used for eligibility traces. This parameter
	 * represents gamma*lambda in a single value. If the value function shouldn't use
	 * traces, set traceDecay to a value smaller than 0.001.
	 */
	public static TabularValueFunction getInstance(double traceDecay) {
		if ( traceDecay < 0.001 )
			return new UnknownStateSpaceFunction();
		return new TracedUnkownStateSpaceFunction(traceDecay);
	}
	
	/**
	 * Use this factory method if the size of the state space is known. Will result in faster
	 * performance than when using {@link #getInstance()}.
	 * 
	 * When using this method, the value function constructs an array of the size of the
	 * state space size, and uses the ID of the states to index into the array. Therefore, the
	 * following points of caution should be observed:
	 * 
	 * 1. The provided state space size can not be exceeded (i.e. the actual state space size
	 * can not be larger than the provided one), or errors might start occurring when using 
	 * the value function.
	 * 
	 * 2. States don't change; more particularly, the state ID uniquely identifies a state. For example,
	 * in a static grid world, a state with a particular ID will always be the same (position
	 * in grid, objects near this position, etc).
	 * 
	 * @param stateSpaceSize the (upper bound on the) size of the state space.
	 * @param traceDecay the decay parameter used for eligibility traces. This parameter
	 * represents gamma*lambda in a single value. If the value function shouldn't use
	 * traces, set traceDecay < 0.001.
	 */
	public static TabularValueFunction getInstance(int stateSpaceSize, double traceDecay) {
		if ( traceDecay < 0.001 )
			return new KnownStateSpaceFunction(stateSpaceSize);
		return new TracedKnownStateSpaceFunction(stateSpaceSize, traceDecay);
	}
	
	/**
	 * gamma*lambda for eligibility traces.
	 */
	protected double traceDecay;
	
	protected TabularValueFunction() {}

	/**
	 * @inheritDoc
	 */
	@Override
	public double getTraceDecay() {
		return traceDecay;
	}

	/**
	 * @inheritDoc
	 */
	@Override
	public void setTraceDecay(double decay) {
		traceDecay = decay;			
	}
	
	protected static class UnknownStateSpaceFunction extends TabularValueFunction {
		
		private static final long serialVersionUID = 1556916165824138600L;
		
		//Performance test showed Trove only outperforms Java HashMap if state space larger than
		//around 250,000 states or so; for smaller spaces, it is actually slower. On the assumption
		//that most uses will involve less states, using HashMap.
		//protected TObjectDoubleHashMap<State> table;
		/**
		 * Table of values for each state.
		 */
		protected Map<State, Double> table;
		
		private State lastState, prevLastState;
		private double lastVal, prevLastVal;
		
		protected UnknownStateSpaceFunction() {
			table = new HashMap<State, Double>(5000);
			//table = new TObjectDoubleHashMap<State>(5000);
		}

		/**
		 * Gets the value of the specified state.
		 * 
		 * @param s the state of which to retrieve the value.
		 * @return the value of the state, or 0 if the state didn't have any value associated yet.
		 * (This method also returns 0 if the state does have an associated value, but that value
		 * is 0.)
		 */
		@Override
		public <T> double getValue(State<T> s) {
			updateMemory(s);
			return lastVal;
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void setValue(State<T> s, double value) {
			table.put(s, value);
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void update(State<T> s, double delta) {
			Double val = table.get(s);
			update(s, (val == null ? 0 : val), delta);
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public void update(double delta) {
			update(lastState, lastVal, delta);		
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public void updatePrevious(double delta) {
			update(prevLastState, prevLastVal, delta);
		}
		
		protected <T> void update(State<T> s, double stateVal, double delta) {
            table.put(s, stateVal + delta);   
		}
		
		/**
		 * Updates the function's "memory" of previous states, and retrieved
		 * values, in order to be able to implement methods {@link #update(double)}
		 * and {@link #updatePrevious(double)}.
		 * 
		 * @param s the state currently being queried
		 */
		protected <T> void updateMemory(State<T> s) {
			prevLastState = lastState;
			lastState = s;
			prevLastVal = lastVal;
			
			Double d = table.get(s);
			if ( d == null )
				lastVal = 0;
			else
				lastVal = d;
		}
		
		public void init() {
			reset();
			table.clear();
		}
		
		public void reset() {
			prevLastState = lastState = null;
			prevLastVal = lastVal = 0;
		}
		
		public String toString() {
			String s = "";
			for ( Map.Entry<State, Double> e : table.entrySet() )
				s += e.getKey() + "=" + e.getValue() + "\n";
			return s;
				
		}
	}
	
	protected static class TracedUnkownStateSpaceFunction extends UnknownStateSpaceFunction {
		
		private static final long serialVersionUID = 7747881685343532998L;
		
		/**
		 * Table of traces for each state. Use an array instead of a single double to also
		 * store the value of the state while it's "tracing"; this avoids having to
		 * retrieve the value from the value table each time the state is updated.
		 * 
		 * First entry in the array is trace, second is value.
		 */
		protected Map<State, double[]> traceTable;
		
		protected TracedUnkownStateSpaceFunction(double traceDecay) {
			traceTable = new HashMap<State, double[]>(100);
			setTraceDecay(traceDecay);
		}
		
		@Override
		protected <T> void update(State<T> s, double stateVal, double delta) {
			//check if the state to update is already in the trace table.
			//if might also not be in the value table yet (if it's the 1st time we encounter
			//this state), but this is taken care of in the update over all
			//tracing states.
			double eAndV[] = traceTable.get(s);
			if ( eAndV == null ) {
				eAndV = new double[] {0, stateVal};
				traceTable.put(s, eAndV);
			}
			eAndV[0] += 1;
			
			updateAll(delta);
		}
		
		/**
		 * Update all states that have traces above threshold 0.001.
		 * 
		 * @param delta the delta to update with.
		 */
		protected void updateAll(double delta) {
			double traceAndVal[];
			List<Map.Entry<State, double[]>> toRemove = new ArrayList<Map.Entry<State, double[]>>();
			Set<Map.Entry<State, double[]>> traceSet = traceTable.entrySet();
			for ( Map.Entry<State, double[]> traceEntry : traceSet ) {
				traceAndVal = traceEntry.getValue();
				
				//update the state value: V(s) = V(s) + delta*e(s)
				traceAndVal[1] += delta*traceAndVal[0];
				table.put(traceEntry.getKey(), traceAndVal[1]);
				
				//decay the trace: e(s) = gamma*lambda*e(s)
				//and remove it if below threshold
				if ( (traceAndVal[0] *= traceDecay) < 0.001 )
					toRemove.add(traceEntry);
			}
			//remove here because we have to wait until iterator is finished
			traceSet.removeAll(toRemove);
		}
	}
	
	protected static class KnownStateSpaceFunction extends TabularValueFunction {
		
		private static final long serialVersionUID = -1545088402932353105L;
		
		protected double table[];
		
		private int lastStateID, prevLastStateID;
		
		protected KnownStateSpaceFunction(int stateSpaceSize) {
			table = new double[stateSpaceSize];
		}
		
		/**
		 * Gets the value of the specified state.
		 * 
		 * @param s the state of which to retrieve the value.
		 * @return the value of the state, or 0 if the state didn't have any value associated yet.
		 * (This method also returns 0 if the state does have an associated value, but that value
		 * is 0.)
		 */
		@Override
		public <T> double getValue(State<T> s) {
			updateMemory(s);
			return table[lastStateID];
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void setValue(State<T> s, double value) {
			table[s.getID()] = value;
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void update(State<T> s, double delta) {
			updateGeneric(s.getID(), delta);
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void update(double delta) {
			updateGeneric(lastStateID, delta);
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public void updatePrevious(double delta) {
			updateGeneric(prevLastStateID, delta);
		}
		
		protected <T> void updateGeneric(int stateID, double delta) {
			table[stateID] += delta;
		}
		
		/**
		 * Updates the function's "memory" of previous states, and retrieved
		 * values, in order to be able to implement methods {@link #update(double)}
		 * and {@link #updatePrevious(double)}.
		 * 
		 * @param s the state currently being queried
		 */
		protected <T> void updateMemory(State<T> s) {
			prevLastStateID = lastStateID;
			lastStateID = s.getID();
		}
		
		public void init() {
			reset();
			table = new double[table.length];
		}
		
		public void reset() {
			prevLastStateID = lastStateID = 0;
		}
		
		public String toString() {
			return Arrays.toString(table);
		}
	}
	
	protected static class TracedKnownStateSpaceFunction extends KnownStateSpaceFunction {
		
		private static final long serialVersionUID = 6410765616678284944L;
		
		/**
		 * Table of traces e(s) for each state.
		 */
		protected Map<Integer, Double> traceTable;
		
		protected TracedKnownStateSpaceFunction(int stateSpaceSize, double traceDecay) {
			super(stateSpaceSize);
			traceTable = new HashMap<Integer, Double>(100);
			setTraceDecay(traceDecay);
		}
		
		@Override
		protected <T> void updateGeneric(int stateID, double delta) {
			Double e = traceTable.get(stateID);
			traceTable.put(stateID, (e == null ? 1 : e+1));
			updateAll(delta);
		}
		
		/**
		 * Update all states that have traces above threshold 0.001.
		 * 
		 * @param delta the delta to update with.
		 */
		protected void updateAll(double delta) {
			double trace;
			List<Map.Entry<Integer, Double>> toRemove = new ArrayList<Map.Entry<Integer, Double>>();
			Set<Map.Entry<Integer, Double>> traceSet = traceTable.entrySet();
			for ( Map.Entry<Integer, Double> traceEntry : traceSet ) {
				trace = traceEntry.getValue();
				table[traceEntry.getKey()] += delta*trace;
				if ( (trace *= traceDecay) < 0.001 )
					toRemove.add(traceEntry);
				else
					traceEntry.setValue(trace);
			}
			traceSet.removeAll(toRemove);
		}
	}
}
